import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image

class ImageClassifier(nn.Module):
    """图像分类器，使用ResNet18作为特征提取器，用于图像分类任务"""
    
    def __init__(self, model_path=None, device=None, num_classes=2, pretrained=True):
        """
        初始化图像分类器
        
        参数:
            model_path (str): 模型权重文件的路径
            device (torch.device): 计算设备
            num_classes (int): 分类的类别数量
            pretrained (bool): 是否使用预训练模型（兼容旧版本参数）
        """
        super(ImageClassifier, self).__init__()
            
        # 初始化基础模型（使用ResNet18，轻量且性能良好）
        # 更新API调用，使用weights参数而非pretrained
        if pretrained:
            self.model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1)
        else:
            self.model = models.resnet18(weights=None)
        
        # 修改最后一层以适应我们的分类任务
        num_ftrs = self.model.fc.in_features
        self.model.fc = nn.Linear(num_ftrs, num_classes)
        
        # 设置设备
        self.device = device if device is not None else torch.device('cpu')
        
        # 如果提供了模型路径，加载模型权重
        if model_path is not None and isinstance(model_path, str):
            self.load_model(model_path)
        
        # 将模型移到指定设备
        self.model = self.model.to(self.device)
        
        # 定义图像预处理
        self.transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        
        # 分类名称
        self.class_names = ['face', 'flower']
    
    def _fix_state_dict_keys(self, state_dict):
        """
        修复状态字典中的键名前缀问题
        
        参数:
            state_dict (dict): 原始状态字典
            
        返回:
            dict: 修复后的状态字典
        """
        fixed_dict = {}
        
        # 检查是否所有键都有"model."前缀
        relevant_keys = [k for k in state_dict.keys() 
                         if k not in ['optimizer_state_dict', 'scheduler_state_dict']]
        if relevant_keys:
            has_model_prefix = all(k.startswith('model.') for k in relevant_keys)
        else:
            has_model_prefix = False
        
        for key, value in state_dict.items():
            if key in ['optimizer_state_dict', 'scheduler_state_dict']:
                continue
                
            if has_model_prefix and key.startswith('model.'):
                # 移除"model."前缀
                new_key = key[6:]  # 'model.' 长度为6
                fixed_dict[new_key] = value
            else:
                # 保持原样
                fixed_dict[key] = value
                
        return fixed_dict
    
    def load_model(self, model_path):
        """
        加载模型权重
        
        参数:
            model_path (str): 模型权重文件的路径
        """
        try:
            # 使用安全模式加载模型
            checkpoint = torch.load(model_path, map_location=self.device, weights_only=True)
            
            # 处理checkpoint可能是字典的情况
            if isinstance(checkpoint, dict):
                if 'model_state_dict' in checkpoint:
                    state_dict = checkpoint['model_state_dict']
                elif 'state_dict' in checkpoint:
                    state_dict = checkpoint['state_dict']
                else:
                    state_dict = checkpoint
            else:
                state_dict = checkpoint
            
            # 修复状态字典键名
            fixed_state_dict = self._fix_state_dict_keys(state_dict)
            
            # 尝试加载修复后的状态字典
            try:
                self.model.load_state_dict(fixed_state_dict)
                print(f"分类器模型已成功加载: {model_path}")
            except Exception as e:
                print(f"加载修复后的状态字典失败: {e}")
                print("尝试进行严格=False的加载...")
                # 尝试非严格加载
                self.model.load_state_dict(fixed_state_dict, strict=False)
                print(f"分类器模型已非严格加载: {model_path}")
                
        except Exception as e:
            print(f"分类器模型加载失败: {e}")
            print("使用预训练的基础分类器")
    
    def forward(self, x):
        """
        前向传播
        
        参数:
            x (torch.Tensor): 输入张量
            
        返回:
            torch.Tensor: 输出张量
        """
        return self.model(x)
    
    def predict_image(self, image, device=None, class_names=None):
        """
        对单个图像进行预测
        
        参数:
            image (PIL.Image): 输入图像
            device (torch.device): 计算设备
            class_names (list): 类别名称列表
            
        返回:
            str: 预测的类别名称
        """
        # 设置设备
        if device is None:
            device = self.device
            
        # 设置类别名称
        if class_names is None:
            class_names = self.class_names
            
        # 确保图像是RGB模式
        if image.mode != 'RGB':
            image = image.convert('RGB')
            
        # 预处理图像
        img_tensor = self.transform(image).unsqueeze(0).to(device)
        
        # 进行推理
        self.eval()
        with torch.no_grad():
            outputs = self(img_tensor)
            _, predicted = torch.max(outputs, 1)
            
        # 返回预测的类别
        return class_names[predicted.item()]
    
    def classify(self, image):
        """
        分类图像的简便方法
        
        参数:
            image (PIL.Image): 输入图像
            
        返回:
            str: 预测的类别名称
        """
        return self.predict_image(image)


# 简单的测试代码
if __name__ == "__main__":
    classifier = ImageClassifier()
    test_image = Image.open("sample_image.jpg")
    result = classifier.predict_image(test_image, torch.device('cpu'), ['face', 'flower'])
    print(f"分类结果: {result}") 